#include "socket.h"

#include<iostream>
using namespace std;
#pragma comment(lib, "ws2_32.lib")

Socket::Socket() 
{
	sck_listen = sck_client = NULL;
}

Socket::~Socket()
{
	CleanUp();
}

void Socket::CleanUp() 
{
	if (sck_listen != NULL || sck_client != NULL) {
		if (sck_listen != NULL)	{
			closesocket(sck_listen);
		}
		if (sck_client != NULL)	{
			closesocket(sck_client);
		}
		WSACleanup();
		sck_listen = sck_client = NULL;
	}
}

bool Socket::Connect(const char* host, const uint32_t port)
{
	if (sck_client == NULL)	{
		if (!InitClientSocket(host, port)) {
			CleanUp();			
			return false;
		}
	}	
	sockaddr_in addr;
	addr.sin_family = AF_INET;
	addr.sin_addr.S_un.S_addr = inet_addr(host);
	addr.sin_port = htons((u_short)port);
	memset(addr.sin_zero, 0, 8);	
	if (connect(sck_client, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR) {	
		cout<<"SOCKET ERROR "<<WSAGetLastError()<<endl;
		CleanUp();	
		return false;
	}	
	return true;
}

void Socket::Wait(const uint32_t port)
{
	if (sck_listen == NULL)	{
		if (!InitServerSocket(port)) {
			CleanUp();
			return;
		}
	}
	if (listen(sck_listen, SOMAXCONN) == SOCKET_ERROR) {
		CleanUp();
		return;
	}	
	sck_client = accept(sck_listen, NULL, NULL);
	if (sck_client == INVALID_SOCKET) {
		CleanUp();
		return;
	}
}

bool Socket::InitServerSocket(const uint32_t port)
{
	WSADATA ws;	
	if (WSAStartup(MAKEWORD(2, 2), &ws) != NO_ERROR) {
		return false;
	}

	char HostName[100] = "";
	gethostname((char *)HostName, sizeof(HostName)-1);	
	struct addrinfo hints, *res, *p;
    int status;
    char ipstr[INET6_ADDRSTRLEN]; 	
    memset(&hints, 0, sizeof hints);
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    if ((status = getaddrinfo(HostName, NULL, &hints, &res)) != 0) {
        fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(status));
        return false;
    }    
    for(p = res;p != NULL; p = p->ai_next) {
        void *addr;            
        if (p->ai_family == AF_INET) {
            struct sockaddr_in *ipv4 = (struct sockaddr_in *)p->ai_addr;
            addr = &(ipv4->sin_addr);			
			inet_ntop(p->ai_family, addr, ipstr, sizeof ipstr);			
			cout<<"Sever ip is "<<ipstr<<endl;
			break;
		}
	}	
    freeaddrinfo(res); 
	
	sck_listen = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (sck_listen == INVALID_SOCKET) {
		return false;
	}
	
	struct sockaddr_in addr;
	addr.sin_family = AF_INET;
	addr.sin_port = htons((u_short)port);
	addr.sin_addr.S_un.S_addr = inet_addr(ipstr);
	memset(addr.sin_zero, 0, 8);
	if (bind(sck_listen, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR) {
		return false;
	}	
	return true;
}

bool Socket::InitClientSocket(const char* host, const uint32_t port)
{
	WSADATA ws;	
	if (WSAStartup(MAKEWORD(2, 2), &ws) != NO_ERROR) {
		return false;
	}	
	sck_client = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (sck_client == INVALID_SOCKET) {
		return false;
	}	
	return true;
}

void Socket::Send(const Serialization content)
{
	if (sck_client == NULL)	{
		return;
	}
	uint32_t size = content.getSize();
	send(sck_client, (char*)&size, 4, 0);	

	uint32_t pos = 0;
	while (pos < content.getSize()) {
		int result = send(sck_client, content.getData()+pos, content.getSize()-pos, 0);
		if (result > 0) {
			pos += result;
		} else {
			cout<<"Error!"<<endl;
			break;
		}
	}    
}

Serialization Socket::Receive()
{
	uint32_t n;
	if (sck_client == NULL)	{
		Serialization s(NULL, 0);
		return s;
	}	
	// Get length first;
	recv(sck_client, (char*)&n, 4, 0);	
	char* p = new char[n];
	uint32_t pos = 0;
	// Get data;	
	while (pos < n) {
		int result = recv(sck_client, p+pos, n-pos, 0);
		if (result > 0) {
			pos += result;
		} else {
			cout<<"Error!"<<endl;
			break;
		}
	}	
	Serialization s(p, n);	
	return s;
}